# -*- coding:utf-8 -*-
# @FileName :pos_tag.py
# @Date: 2024/7/31
# @Author:天空之城
# 地址词性识别
from address_normalizer.app.models.base import BaseModel


class PosTag(BaseModel):
    def __init__(self):
        super().__init__()
        self.tag_map = {
            'province': 'province',
            'city': 'city',
            'district': 'district',
        }

    def rule_tagger(self, sentence):
        self.processor = self.init_nlp()
        doc = self.processor(sentence)
        return {self.tag_map.get(span.label_, 'other_address'): span.text for span in doc.ents}

    def jieba_tagger(self, sentence):
        result = {}
        self.processor = self.init_jieba()
        words = self.processor.cut(sentence)
        for word, flag in words:
            if flag in self.tag_map and self.tag_map[flag] not in result:
                result[self.tag_map[flag]] = word
        return result



